#include "gtd_common.hpp"
#include "gtd_tracker.hpp"
#include "gtd_profiler.hpp"
#include "gtd_convolution.hpp"

NAMESPACE_BEG(gtd)

FeatureTracker::FeatureTracker(uint width, uint height, uint targetCount,
                               uint minDist, uint levels, uint subsample): 
width(width), 
height(height),
npixels(width * height),
targetCount(targetCount),
minDist(minDist),
levels(levels),
subsample(subsample),
eigenvalues(npixels),
distMap(((npixels - 1) >> 5) + 1),
lostMinDet(0),
lostMaxIts(0),
lostMaxRes(0),
old(0),
cur(1)
{
    computeGaussianKernels(0.7f, smoothKernels);
    computeGaussianKernels(0.5f * subsample, pyramidKernels);
    computeGaussianKernels(1.0f, gradientKernels);

    uint w = width;
    uint h = height;
    for(uint i = 0; i < levels; ++i)
    {
        uint2 dims = { w, h };
        imageDims.push_back(dims);
        imageSizes.push_back(w * h * sizeof(float));
        for(uint j = 0; j < IMAGE_TYPE_COUNT; ++j)
        {
            for(uint k = 0; k < 2; ++k)
            {
                float* mem;
                CUDA_ERROR_CHECK(cudaMalloc(&mem, imageSizes.back()));
                images[j][k].push_back(mem);
            }
        }
        w /= subsample;
        h /= subsample;
    }

    CUDA_ERROR_CHECK(cudaMalloc(&d_eigenvalues, imageSizes[0]));
    CUDA_ERROR_CHECK(cudaMalloc(&d_features, targetCount * sizeof(FeaturePair)));
}

FeatureTracker::~FeatureTracker()
{
    CUDA_ERROR_CHECK(cudaFree(d_features));
    CUDA_ERROR_CHECK(cudaFree(d_eigenvalues));

    for(uint i = 0; i < IMAGE_TYPE_COUNT; ++i)
    {
        for(uint j = 0; j < 2; ++j)
        {
            for(uint k = 0; k < levels; ++k)
            {
                CUDA_ERROR_CHECK(cudaFree(images[i][j][k]));
            }
        }
    }

    std::cout << "Lost Min Det: " << lostMinDet << std::endl;
    std::cout << "Lost Max Its: " << lostMaxIts << std::endl;
    std::cout << "Lost Max Res: " << lostMaxRes << std::endl;
}

void FeatureTracker::update(std::vector<FeaturePair>& pairs, 
                            const uchar4* d_frame)
{
    PROFILE("Feature Tracking");

    // swap image resources from last frame
    old ^= 1;
    cur ^= 1;

    // generate smoothed gaussian pyramid from new frame data
    preprocess(d_frame);

    // track features between current and last frame
    pairs.clear();
    if(!features.empty())
    {
        trackOldFeatures();
        pairs = features;
    }

    // select new features if some were lost
    if(features.size() < targetCount)
        selectNewFeatures();
}

const std::vector<float*>& FeatureTracker::getPyramid(ImageType type) const
{
    return images[type][cur];
}

const std::vector<uint2>& FeatureTracker::getDims() const
{
    return imageDims;
}

void FeatureTracker::preprocess(const uchar4* d_frame)
{
    PROFILE("Preprocessing");

    {
        PROFILE("Intensity Normalization");
        
        // convert to intensity and normalize
        normalizeImage(images[TEMP][0][0], d_frame, imageDims[0]);

        PROFILER_CUDA_BARRIER();
    }

    {
        PROFILE("Smoothing");

        // smooth new intensities
        convolveHorizontal(images[TEMP][1][0], images[TEMP][0][0], 
            width, height, smoothKernels[0]);
        convolveVertical(images[INTENSITY][cur][0], images[TEMP][1][0],
            width, height, smoothKernels[0]);

        PROFILER_CUDA_BARRIER();
    }

    // compute gaussian pyramid
    for(uint i = 0; i < levels; ++i)
    {
        PROFILE("Pyramid");

        // compute level gradients
        convolveVertical(images[TEMP][0][i], images[INTENSITY][cur][i],
            imageDims[i].x, imageDims[i].y, gradientKernels[0]);
        convolveHorizontal(images[TEMP][1][i], images[INTENSITY][cur][i],
            imageDims[i].x, imageDims[i].y, gradientKernels[0]);
        convolveHorizontal(images[GRADX][cur][i], images[TEMP][0][i],
            imageDims[i].x, imageDims[i].y, gradientKernels[1]);
        convolveVertical(images[GRADY][cur][i], images[TEMP][1][i],
            imageDims[i].x, imageDims[i].y, gradientKernels[1]);

        PROFILER_CUDA_BARRIER();

        // downsample if not last level
        if(i < (levels - 1))
        {
            PROFILE("Downsampling");

            // smooth again before downsampling
            convolveHorizontal(images[TEMP][0][i], images[INTENSITY][cur][i],
                imageDims[i].x, imageDims[i].y, pyramidKernels[0]);
            convolveVertical(images[TEMP][1][i], images[TEMP][0][i],
                imageDims[i].x, imageDims[i].y, pyramidKernels[0]);

            // run downsampling kernel
            downsample(images[INTENSITY][cur][i+1], imageDims[i+1], 
                images[TEMP][1][i], imageDims[i]);

            PROFILER_CUDA_BARRIER();
        }

        PROFILER_CUDA_BARRIER();
    }

    PROFILER_CUDA_BARRIER();
}

void FeatureTracker::trackOldFeatures()
{
    PROFILE("Tracking");

    {
        PROFILE("Reset coordinates");

        for(uint i = 0; i < features.size(); ++i)
        {
            features[i].oldPos = features[i].newPos;
        }
    }
    
    {
        PROFILE("Copy features to GPU");

        CUDA_ERROR_CHECK(cudaMemcpy(d_features, &features[0], 
            features.size() * sizeof(FeaturePair), cudaMemcpyHostToDevice));
    }

    {
        PROFILE("Pyramid computation");

        float sub = 1.0f / std::pow((float)subsample, (int)levels - 1);

        for(int i = levels - 1; i >= 0; --i)
        {
            trackFeatures(d_features, features.size(),
                images[INTENSITY][old][i], images[INTENSITY][cur][i],
                images[GRADX][old][i], images[GRADX][cur][i],
                images[GRADY][old][i], images[GRADY][cur][i],
                imageDims[i], sub);

            sub = (float)subsample;
        }

        CUDA_ERROR_CHECK(cudaThreadSynchronize());
    }

    {
        PROFILE("Copy features from GPU");

        // copy back feature pairs
        CUDA_ERROR_CHECK(cudaMemcpy(&features[0], d_features,
            features.size() * sizeof(FeaturePair), cudaMemcpyDeviceToHost));

        // finalize feature pairs that were successfully tracked
        for(size_t i = 0; i < features.size();)
        {
            if((features[i].newPos.x >= minDist) && (features[i].newPos.x < (width - minDist)) &&
               (features[i].newPos.y >= minDist) && (features[i].newPos.y < (height - minDist)))
            {
                ++i;
            }
            else
            {
                if(features[i].newPos.x == -2.0f)
                {
                    ++lostMinDet;
                }
                else if(features[i].newPos.x == -4.0f)
                {
                    ++lostMaxIts;
                }
                else if(features[i].newPos.x == -3.0f)
                {
                    ++lostMaxRes;
                }

                features[i] = features.back();
                features.pop_back();
            }
        }
    }

    //std::cout << "Tracked: " << features.size() << std::endl;
}

void FeatureTracker::selectNewFeatures()
{
    PROFILE("Selection");

    {
        PROFILE("Eigenvalues");

        minEigenvalue(d_eigenvalues, images[GRADX][cur][0],
            images[GRADY][cur][0], imageDims[0]);

        CUDA_ERROR_CHECK(cudaThreadSynchronize());
    }

    {
        PROFILE("Copy eigenvalues from GPU");

        CUDA_ERROR_CHECK(cudaMemcpy(&eigenvalues[0], d_eigenvalues, 
            imageSizes[0], cudaMemcpyDeviceToHost));
    }

    {
        PROFILE("Eigenvalue Thresholding");

        eigenpoints.clear();
        for(uint i = minDist; i < height - minDist; ++i)
        {
            for(uint j = minDist; j < width - minDist; ++j)
            {
                if(eigenvalues[i * width + j] > 0.01f)
                {
                    EigenPoint point = { eigenvalues[i * width + j], { j, i } };
                    eigenpoints.push_back(point);
                }
            }
        }
    }

    {
        PROFILE("Eigenvalue Heapify");

        std::make_heap(eigenpoints.begin(), eigenpoints.end());
    }

    {
        PROFILE("Enforce distance");

        // fill the feature map with previously selected features
        std::memset(&distMap[0], 0, distMap.size() * sizeof(Uint32));
        for(size_t i = 0; i < features.size(); ++i)
        {
            uint maxX = std::min((uint)features[i].newPos.x + minDist, width);
            uint maxY = std::min((uint)features[i].newPos.y + minDist, height);

            for(uint y = std::max((int)features[i].newPos.y - (int)minDist, 0); y < maxY; ++y)
            {
                for(uint x = std::max((int)features[i].newPos.x - (int)minDist, 0); x < maxX; ++x)
                {
                    uint p = y * width + x;
                    distMap[p >> 5] |= 1 << (p & 31);
                }
            }
        }

        // fill remianing slots with newly found trackable features
        while(!eigenpoints.empty() && (features.size() < targetCount))
        {
            uint p = eigenpoints[0].pos.y * width + eigenpoints[0].pos.x;
            
            if((distMap[p >> 5] & (1 << (p & 31))) == 0)
            {
                uint maxX = std::min(eigenpoints[0].pos.x + minDist, width);
                uint maxY = std::min(eigenpoints[0].pos.y + minDist, height);
                
                for(uint y = std::max((int)eigenpoints[0].pos.y - (int)minDist, 0); y < maxY; ++y)
                {
                    for(uint x = std::max((int)eigenpoints[0].pos.x - (int)minDist, 0); x < maxX; ++x)
                    {
                        uint p2 = y * width + x;
                        distMap[p2 >> 5] |= 1 << (p2 & 31);
                    }
                }
                
                FeaturePair pair = 
                { 
                    { eigenpoints[0].pos.x, eigenpoints[0].pos.y },
                    { eigenpoints[0].pos.x, eigenpoints[0].pos.y }
                };
                features.push_back(pair);
            }

            std::pop_heap(eigenpoints.begin(), eigenpoints.end());
            eigenpoints.pop_back();
        }
    }
}

NAMESPACE_END(gtd)
